#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import argparse
import json
import pdfplumber
from PIL import Image, ImageDraw

# 默认红色边框RGB值
DEFAULT_RED_RGB = (237, 2, 140)
DEFAULT_COLOR_TOLERANCE = 40

class PDFFinalMultiRectangleExtractor:
    def __init__(self, color_tolerance=DEFAULT_COLOR_TOLERANCE, target_rgb=DEFAULT_RED_RGB):
        self.color_tolerance = color_tolerance
        self.target_rgb = target_rgb
        # 计算颜色容差范围
        self.red_lower = max(0, target_rgb[0] - color_tolerance)
        self.red_upper = min(255, target_rgb[0] + color_tolerance)
        self.green_lower = max(0, target_rgb[1] - color_tolerance)
        self.green_upper = min(255, target_rgb[1] + color_tolerance)
        self.blue_lower = max(0, target_rgb[2] - color_tolerance)
        self.blue_upper = min(255, target_rgb[2] + color_tolerance)
        
        print(f"使用目标RGB值: {target_rgb}，容差: {color_tolerance}")
        print(f"颜色范围 - R: [{self.red_lower}-{self.red_upper}], G: [{self.green_lower}-{self.green_upper}], B: [{self.blue_lower}-{self.blue_upper}]")
    
    def is_red_color(self, rgb):
        """检查颜色是否在目标红色范围内"""
        if len(rgb) != 3:
            return False
        
        r, g, b = rgb
        return (self.red_lower <= r <= self.red_upper and
                self.green_lower <= g <= self.green_upper and
                self.blue_lower <= b <= self.blue_upper)
    
    def find_all_red_pixels(self, img):
        """找出图像中所有的红色像素"""
        width, height = img.size
        red_pixels = []
        
        for y in range(0, height, 1):
            for x in range(0, width, 1):
                pixel = img.getpixel((x, y))
                if self.is_red_color(pixel):
                    red_pixels.append((x, y))
        
        return red_pixels
    
    def is_point_in_rectangle(self, point, rect):
        """检查点是否在矩形内"""
        x, y = point
        x1, y1, x2, y2 = rect['x1'], rect['y1'], rect['x2'], rect['y2']
        return x1 <= x <= x2 and y1 <= y <= y2
    
    def distance_between_points(self, p1, p2):
        """计算两点之间的距离"""
        return ((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** 0.5
    
    def find_independent_rectangles(self, red_pixels, width, height, min_size=50, max_gap=10):
        """从红色像素中寻找多个独立的矩形"""
        if not red_pixels:
            return []
        
        print(f"开始寻找独立矩形，红色像素总数: {len(red_pixels)}")
        
        # 首先创建红色像素的二维数组表示，用于快速查找
        pixel_matrix = [[False for _ in range(width)] for _ in range(height)]
        for x, y in red_pixels:
            if 0 <= x < width and 0 <= y < height:
                pixel_matrix[y][x] = True
        
        # 存储已访问的像素
        visited = set()
        
        # 存储检测到的矩形
        rectangles = []
        
        # 定义8个方向的邻居
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0),
                      (1, 1), (1, -1), (-1, -1), (-1, 1)]
        
        # 扫描图像中的红色像素
        for x, y in red_pixels:
            if (x, y) not in visited:
                # 使用BFS找到连通区域
                queue = [(x, y)]
                visited.add((x, y))
                
                # 初始化连通区域的边界
                min_x, min_y = x, y
                max_x, max_y = x, y
                
                # 用于存储连通区域的所有点
                region_points = [(x, y)]
                
                while queue:
                    curr_x, curr_y = queue.pop(0)
                    
                    # 更新区域边界
                    min_x = min(min_x, curr_x)
                    max_x = max(max_x, curr_x)
                    min_y = min(min_y, curr_y)
                    max_y = max(max_y, curr_y)
                    
                    # 检查所有邻居
                    for dx, dy in directions:
                        nx, ny = curr_x + dx, curr_y + dy
                        # 确保在图像范围内且未被访问且是红色像素
                        if (0 <= nx < width and 0 <= ny < height and 
                            (nx, ny) not in visited and pixel_matrix[ny][nx]):
                            queue.append((nx, ny))
                            visited.add((nx, ny))
                            region_points.append((nx, ny))
                    
                # 计算矩形的宽度和高度
                rect_width = max_x - min_x
                rect_height = max_y - min_y
                
                # 只保留足够大的矩形
                if rect_width >= min_size and rect_height >= min_size:
                    # 简单判断是否为矩形：检查边界上的点数量
                    border_points = 0
                    left_edge_points = 0  # 左边框上的点
                    right_edge_points = 0  # 右边框上的点
                    top_edge_points = 0   # 上边框上的点
                    bottom_edge_points = 0  # 下边框上的点
                    
                    for px, py in region_points:
                        # 检查点是否在矩形边界附近
                        if abs(px - min_x) <= max_gap:
                            left_edge_points += 1
                        elif abs(px - max_x) <= max_gap:
                            right_edge_points += 1
                        elif abs(py - min_y) <= max_gap:
                            top_edge_points += 1
                        elif abs(py - max_y) <= max_gap:
                            bottom_edge_points += 1
                        
                    border_points = left_edge_points + right_edge_points + top_edge_points + bottom_edge_points
                    
                    # 计算边界点比例和填充率
                    border_ratio = border_points / len(region_points) if region_points else 0
                    area_ratio = len(region_points) / (rect_width * rect_height) if (rect_width * rect_height) > 0 else 0
                    
                    # 计算长宽比
                    aspect_ratio = max(rect_width, rect_height) / min(rect_width, rect_height) if min(rect_width, rect_height) > 0 else 0
                    
                    # 检查四个边是否都有足够的点（边框完整性）
                    has_complete_border = (left_edge_points > rect_height * 0.3 and 
                                          right_edge_points > rect_height * 0.3 and 
                                          top_edge_points > rect_width * 0.3 and 
                                          bottom_edge_points > rect_width * 0.3)
                    
                    # 对于边框线形成的矩形，边界点比例会较高，填充率会较低，且有较完整的边界
                    is_rectangle = ((border_ratio > 0.6 and area_ratio < 0.25 and 
                                   aspect_ratio < 10 and has_complete_border) or 
                                  (rect_width * rect_height > 5000 and border_ratio > 0.5 and 
                                   area_ratio < 0.4 and has_complete_border))
                    
                    print(f"发现连通区域: 尺寸={rect_width}x{rect_height}, 点数={len(region_points)}, 边界比例={border_ratio:.2f}, "
                          f"填充率={area_ratio:.2f}, 长宽比={aspect_ratio:.2f}, 边界完整性={has_complete_border}")
                    
                    if is_rectangle:
                        rectangles.append({
                            'x1': min_x,
                            'y1': min_y,
                            'x2': max_x,
                            'y2': max_y,
                            'type': 'independent_rectangle',
                            'point_count': len(region_points),
                            'border_ratio': border_ratio,
                            'area_ratio': area_ratio
                        })
        
        # 对矩形按面积排序，大的在前
        rectangles.sort(key=lambda r: (r['x2'] - r['x1']) * (r['y2'] - r['y1']), reverse=True)
        
        print(f"初步检测到 {len(rectangles)} 个可能的矩形")
        
        # 去重嵌套矩形（保留最外层或最大的矩形）
        unique_rectangles = []
        for rect in rectangles:
            is_nested = False
            for existing_rect in unique_rectangles:
                # 检查当前矩形是否完全包含在已存在的矩形中（允许小的容差）
                if (rect['x1'] >= existing_rect['x1'] - max_gap and 
                    rect['y1'] >= existing_rect['y1'] - max_gap and 
                    rect['x2'] <= existing_rect['x2'] + max_gap and 
                    rect['y2'] <= existing_rect['y2'] + max_gap):
                    is_nested = True
                    break
            if not is_nested:
                unique_rectangles.append(rect)
        
        print(f"去重后检测到 {len(unique_rectangles)} 个独立矩形")
        return unique_rectangles
    
    def extract_content_from_rectangle(self, full_page_img, rect, output_dir, page_num, rect_idx):
        """从矩形中提取内容并保存为图像"""
        try:
            # 提取矩形内的内容
            x1, y1, x2, y2 = rect['x1'], rect['y1'], rect['x2'], rect['y2']
            
            # 确保坐标有效
            if x1 >= x2 or y1 >= y2:
                print(f"矩形 {rect_idx+1} 坐标无效: ({x1}, {y1}, {x2}, {y2})")
                return None
            
            # 提取矩形内容
            rect_content = full_page_img.crop((x1, y1, x2, y2))
            
            # 保存矩形内容图像
            img_path = os.path.join(output_dir, f"page_{page_num}_rectangle_{rect_idx+1}.png")
            rect_content.save(img_path)
            
            print(f"成功保存矩形 {rect_idx+1} 内容到 {img_path}")
            return img_path
        except Exception as e:
            print(f"提取矩形 {rect_idx+1} 内容时出错: {str(e)}")
            return None
    
    def save_annotated_mask(self, mask_img, rectangles, output_dir, page_num):
        """保存带有矩形标注的掩码图像"""
        try:
            # 创建可绘制的副本
            annotated_mask = mask_img.copy()
            draw = ImageDraw.Draw(annotated_mask)
            
            # 在掩码上绘制检测到的矩形
            colors = [(255, 255, 0), (0, 255, 0), (0, 255, 255), (255, 0, 255)]  # 黄色、绿色、青色、紫色
            for idx, rect in enumerate(rectangles):
                x1, y1, x2, y2 = rect['x1'], rect['y1'], rect['x2'], rect['y2']
                # 循环使用不同颜色
                color = colors[idx % len(colors)]
                draw.rectangle([(x1, y1), (x2, y2)], outline=color, width=3)
                # 标记矩形序号
                draw.text((x1 + 10, y1 + 10), f"{idx+1}", fill=color)
            
            # 保存标注后的掩码图像
            img_path = os.path.join(output_dir, f"page_{page_num}_annotated_mask.png")
            annotated_mask.save(img_path)
            
            print(f"成功保存标注后的掩码图像到 {img_path}")
            return img_path
        except Exception as e:
            print(f"保存标注后的掩码图像时出错: {str(e)}")
            return None
    
    def process_page(self, pdf_path, page_num, output_dir):
        """处理PDF页面，提取所有独立的红色矩形内容"""
        page_results = {
            "page_number": page_num,
            "detected_rectangles": [],
            "extracted_images": []
        }
        
        try:
            # 使用pdfplumber打开PDF文件
            with pdfplumber.open(pdf_path) as pdf:
                # 获取页面
                if page_num > len(pdf.pages):
                    print(f"错误: 页码 {page_num} 超出PDF总页数 {len(pdf.pages)}")
                    return page_results
                
                page = pdf.pages[page_num - 1]  # pdfplumber页面索引从0开始
                
                # 将页面转换为高分辨率图像
                img = page.to_image(resolution=300)
                full_page_img = img.original
                width, height = full_page_img.size
                
                print(f"处理页面 {page_num}，图像尺寸: {width}x{height}")
                
                # 创建红色掩码图像
                mask_img = Image.new("RGB", (width, height), (0, 0, 0))
                draw = ImageDraw.Draw(mask_img)
                
                # 扫描图像，标记红色像素
                red_pixels = []
                for y in range(0, height, 1):
                    for x in range(0, width, 1):
                        try:
                            pixel = full_page_img.getpixel((x, y))
                            if self.is_red_color(pixel):
                                red_pixels.append((x, y))
                                draw.point((x, y), fill=(255, 0, 0))
                        except Exception:
                            continue
                
                # 保存原始红色掩码图像
                original_mask_path = os.path.join(output_dir, f"page_{page_num}_original_mask.png")
                mask_img.save(original_mask_path)
                print(f"已保存原始红色掩码图像到 {original_mask_path}")
                print(f"页面 {page_num}: 检测到 {len(red_pixels)} 个红色像素")
                
                # 尝试从红色像素中检测多个独立矩形
                detected_rectangles = self.find_independent_rectangles(red_pixels, width, height)
                
                # 如果没有检测到独立矩形，使用简单的边界框方法
                if not detected_rectangles and red_pixels:
                    print("未检测到独立矩形，使用边界框方法")
                    # 按x坐标排序
                    sorted_by_x = sorted(red_pixels, key=lambda p: p[0])
                    # 按y坐标排序
                    sorted_by_y = sorted(red_pixels, key=lambda p: p[1])
                    
                    min_x = sorted_by_x[0][0]
                    max_x = sorted_by_x[-1][0]
                    min_y = sorted_by_y[0][1]
                    max_y = sorted_by_y[-1][1]
                    
                    # 计算宽度和高度
                    width_box = max_x - min_x
                    height_box = max_y - min_y
                    
                    if width_box >= 50 and height_box >= 50:
                        detected_rectangles = [{
                            'x1': min_x,
                            'y1': min_y,
                            'x2': max_x,
                            'y2': max_y,
                            'type': 'bounding_box',
                            'estimated_pixel_count': len(red_pixels)
                        }]
                
                print(f"页面 {page_num}: 检测到 {len(detected_rectangles)} 个矩形")
                
                # 保存带有矩形标注的掩码图像
                annotated_mask_path = self.save_annotated_mask(mask_img, detected_rectangles, output_dir, page_num)
                if annotated_mask_path:
                    page_results["annotated_mask_path"] = annotated_mask_path
                
                # 保存整页图像
                full_page_path = os.path.join(output_dir, f"page_{page_num}_full_page.png")
                full_page_img.save(full_page_path)
                print(f"已保存整页图像到 {full_page_path}")
                
                # 从每个矩形中提取内容
                extracted_images = 0
                for rect_idx, rect in enumerate(detected_rectangles):
                    img_path = self.extract_content_from_rectangle(
                        full_page_img, rect, output_dir, page_num, rect_idx
                    )
                    if img_path:
                        extracted_images += 1
                        page_results["extracted_images"].append({
                            "rectangle_index": rect_idx,
                            "path": img_path,
                            "coordinates": {
                                "x1": rect['x1'],
                                "y1": rect['y1'],
                                "x2": rect['x2'],
                                "y2": rect['y2']
                            },
                            "dimensions": {
                                "width": rect['x2'] - rect['x1'],
                                "height": rect['y2'] - rect['y1'],
                                "area": (rect['x2'] - rect['x1']) * (rect['y2'] - rect['y1'])
                            }
                        })
                
                # 为每个矩形添加尺寸信息
                for rect in detected_rectangles:
                    rect["width"] = rect['x2'] - rect['x1']
                    rect["height"] = rect['y2'] - rect['y1']
                    rect["area"] = rect["width"] * rect["height"]
                
                # 更新页面结果
                page_results["detected_rectangles"] = detected_rectangles
                page_results["red_pixel_count"] = len(red_pixels)
                
                print(f"页面 {page_num}: 成功提取 {extracted_images} 个矩形内容")
        except Exception as e:
            print(f"处理页面 {page_num} 时出错: {str(e)}")
        
        return page_results
    
    def process_pdf(self, pdf_path, output_dir, page_numbers=None):
        """处理PDF文件并提取所有页面中的红色矩形内容"""
        results = {
            "pdf_filename": os.path.basename(pdf_path),
            "pages": []
        }
        
        # 确保输出目录存在
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            print(f"创建输出目录: {output_dir}")
        
        try:
            # 使用pdfplumber打开PDF文件
            with pdfplumber.open(pdf_path) as pdf:
                total_pages = len(pdf.pages)
                print(f"成功打开PDF文件: {pdf_path}，共 {total_pages} 页")
                
                # 确定要处理的页码
                if page_numbers is None:
                    pages_to_process = range(1, total_pages + 1)
                else:
                    pages_to_process = page_numbers
                
                # 处理指定的页面
                for page_num in pages_to_process:
                    if 1 <= page_num <= total_pages:
                        print(f"开始处理页面 {page_num}/{total_pages}")
                        page_result = self.process_page(pdf_path, page_num, output_dir)
                        results["pages"].append(page_result)
                    else:
                        print(f"警告: 页码 {page_num} 超出范围，跳过")
        except Exception as e:
            print(f"处理PDF文件时出错: {str(e)}")
        
        return results

def main():
    parser = argparse.ArgumentParser(description="最终版PDF多矩形红色边框内容提取工具")
    parser.add_argument("--pdf-path", required=True, help="PDF文件路径")
    parser.add_argument("--output-dir", default=None, 
                        help="提取内容的输出目录 (默认根据PDF文件名自动生成)")
    parser.add_argument("--color-tolerance", type=int, default=DEFAULT_COLOR_TOLERANCE,
                        help="颜色检测的容差度")
    parser.add_argument("--target-rgb", type=str, default=f"{DEFAULT_RED_RGB[0]},{DEFAULT_RED_RGB[1]},{DEFAULT_RED_RGB[2]}",
                        help="目标颜色的RGB值，格式为'r,g,b'，例如 '238,29,35'")
    parser.add_argument("--pages", type=str, default="all",
                        help="要处理的页码，如 '1,3,5' 或 'all'")
    args = parser.parse_args()
    
    # 检查PDF文件是否存在
    if not os.path.exists(args.pdf_path):
        print(f"错误: PDF文件不存在: {args.pdf_path}")
        return
    
    # 根据PDF文件名创建子目录
    pdf_filename = os.path.basename(args.pdf_path)
    pdf_name_without_ext = os.path.splitext(pdf_filename)[0]
    # 输出目录设置为images目录下以PDF文件名命名的子目录
    args.output_dir = os.path.join("images", pdf_name_without_ext)
    
    # 解析要处理的页码
    page_numbers = None
    if args.pages.lower() != "all":
        try:
            page_numbers = [int(p) for p in args.pages.split(",")]
        except ValueError:
            print(f"错误: 无效的页码格式: {args.pages}")
            return
    
    # 解析目标RGB值
    try:
        target_rgb = tuple(map(int, args.target_rgb.split(',')))
        if len(target_rgb) != 3 or not all(0 <= c <= 255 for c in target_rgb):
            raise ValueError
    except ValueError:
        print(f"错误: 无效的RGB格式: {args.target_rgb}，使用默认值 {DEFAULT_RED_RGB}")
        target_rgb = DEFAULT_RED_RGB
    
    # 创建提取器实例
    extractor = PDFFinalMultiRectangleExtractor(color_tolerance=args.color_tolerance, target_rgb=target_rgb)
    
    # 处理PDF
    print(f"开始处理PDF文件: {args.pdf_path}")
    results = extractor.process_pdf(args.pdf_path, args.output_dir, page_numbers)
    
    # 确保json_file目录存在
    if not os.path.exists("json_file"):
        os.makedirs("json_file")
        print("创建json_file目录")
    
    # 确保images目录存在
    if not os.path.exists("images"):
        os.makedirs("images")
        print("创建images目录")
    
    # 保存结果到JSON文件（根据PDF文件名命名）
    pdf_filename = os.path.basename(args.pdf_path)
    pdf_name_without_ext = os.path.splitext(pdf_filename)[0]
    json_path = os.path.join("json_file", f"{pdf_name_without_ext}_rectangle_results.json")
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    
    print(f"处理完成，结果已保存到: {json_path}")
    
    # 统计信息
    total_rectangles = sum(len(page["detected_rectangles"]) for page in results["pages"])
    total_images = sum(len(page["extracted_images"]) for page in results["pages"])
    print(f"总共检测到 {total_rectangles} 个红色矩形，提取了 {total_images} 张图像")

if __name__ == "__main__":
    main()